#include "StdAfx.h"
#include "DllContextMenu.h"

const CComBSTRResource ProjectName(IDS_PROJNAME);
const LPCTSTR          CDllContextMenu::RegPath = _T("SOFTWARE\\jachymko\\PEContextMenu");

STDMETHODIMP CDllContextMenu::Initialize(LPCITEMIDLIST, LPDATAOBJECT pDataObj, HKEY)
{
	FORMATETC format  = { CF_HDROP, NULL, DVASPECT_CONTENT, -1, TYMED_HGLOBAL };
	STGMEDIUM storage =	{ TYMED_HGLOBAL };
	
	if (FAILED(pDataObj->GetData(&format, &storage)))
	{
		return E_INVALIDARG;
	}

	HDROP hDrop = static_cast<HDROP>(::GlobalLock(storage.hGlobal));

	if (hDrop == NULL)
	{
		return E_INVALIDARG;
	}

	UINT uNumFiles = ::DragQueryFile(hDrop, 0xFFFFFFFF, NULL, 0);
	TCHAR buff[MAX_PATH];

	for (UINT i = 0; i < uNumFiles; ++i)
	{
		UINT cch = ::DragQueryFile(hDrop, i, buff, MAX_PATH);

		m_files.push_back(CSelectedFile(tstring(buff, cch)));
	}

	::GlobalUnlock(storage.hGlobal);
	::ReleaseStgMedium(&storage);

	return (uNumFiles > 0) ? S_OK : E_INVALIDARG;
}

STDMETHODIMP CDllContextMenu::QueryContextMenu(HMENU hMenu, UINT uMenuIndex, UINT uidFirstCmd, UINT uidLastCmd, UINT uFlags)
{
	if (uFlags & CMF_DEFAULTONLY)
	{
		return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, 0);
	}

	EnsureVerbsLoaded();

	unsigned i = 0;
	unsigned ubound = m_verbs->size() - 1;

	for (verb_iterator it = m_verbs->rbegin(); it != m_verbs->rend(); ++i, ++it)
	{
		if (uidFirstCmd + i > uidLastCmd)
		{
			break;
		}
		
		if (it->CanInvoke(m_files))
		{
			LPCTSTR displayName = it->DisplayName().c_str();
			UINT uidCmd = uidFirstCmd + (ubound - i);

			if (!InsertMenu(hMenu, uMenuIndex, MF_BYPOSITION, uidCmd, displayName))
			{
				return AtlHresultFromLastError();
			}
		}
	}

	return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, i);
}


STDMETHODIMP CDllContextMenu::InvokeCommand(LPCMINVOKECOMMANDINFO pCmdInfo)
{
	if (HIWORD(pCmdInfo->lpVerb) != 0)
	{
		return E_INVALIDARG;
	}

	unsigned int cmd = LOWORD(pCmdInfo->lpVerb);

	if (cmd < m_verbs->size())
	{
		return m_verbs->at(cmd).Invoke(m_files);
	}

	return E_INVALIDARG;
}

STDMETHODIMP CDllContextMenu::GetCommandString(UINT idCmd, UINT uFlags, LPUINT, LPSTR pszName, UINT cchMax)
{
	EnsureVerbsLoaded();

	if (idCmd < m_verbs->size())
	{
		if (uFlags & GCS_HELPTEXT)
		{
			return CopyCommandString(pszName, cchMax, (uFlags & GCS_UNICODE), m_verbs->at(idCmd).HelpString());
		}
	}

	return E_INVALIDARG;
}

STDMETHODIMP CDllContextMenu::CopyCommandString(LPVOID pDest, UINT cchMax, BOOL fUnicode, const tstring& sz)
{
	USES_CONVERSION;

	if (fUnicode)
	{
		return ::StringCchCopyW(reinterpret_cast<LPWSTR>(pDest), cchMax, T2CW(sz.c_str()));
	}

	return ::StringCchCopyA(reinterpret_cast<LPSTR>(pDest), cchMax, T2CA(sz.c_str()));
}
void CDllContextMenu::EnsureVerbsLoaded()
{
	if (m_verbs == NULL)
	{
		m_verbs = new verb_vector();

		CRegKey root;

		if (root.Open(HKEY_LOCAL_MACHINE, RegPath, KEY_READ) == ERROR_SUCCESS)
		{
			DWORD index = 0;
			LRESULT result = ERROR_SUCCESS;

			do
			{
				TCHAR szName[32768] = { 0 };
				DWORD cchName = _countof(szName);

				result = root.EnumKey(index, szName, &cchName);

				if (result == ERROR_SUCCESS)
				{
					CRegKey subkey;
					tstring name(szName, cchName);

					if (subkey.Open(root, name.c_str(), KEY_READ) == ERROR_SUCCESS)
					{
						m_verbs->push_back(CVerb(name, subkey));
					}
				}

				++index;
			}
			while (result == ERROR_SUCCESS);
		}
	}
}

CDllContextMenu::CVerb::CVerb(const tstring& name, CRegKey& key)
	: m_Name(name)
	, m_Type(FileTypeAny)
	, m_DisplayName(QueryStringValue(key, NULL))
	, m_Template(QueryStringValue(key, _T("Command")))
	, m_HelpString(QueryStringValue(key, _T("HelpString")))
{
	DWORD type;

	if (key.QueryDWORDValue(_T("Type"), type) == ERROR_SUCCESS)
	{
		m_Type = type;
	}
}

bool CDllContextMenu::CVerb::CanInvoke(const CSelectedFile& item) const
{
	return (item.Type() & m_Type) != 0;
}

bool CDllContextMenu::CVerb::CanInvoke(const item_vector& items) const
{
	for (item_iterator it = items.begin(); it != items.end(); ++it)
	{
		if (CanInvoke(*it))
		{
			return true;
		}
	}

	return false;
}

HRESULT CDllContextMenu::CVerb::Invoke(const item_vector& items) const
{
	vector<tstring> paths;

	for (item_iterator it = items.begin(); it != items.end(); ++it)
	{
		if (CanInvoke(*it))
		{
			paths.push_back(it->Path());
		}
	}

	if (!paths.empty())
	{
		tstring cmdline = m_Template.Process(paths);

		HRESULT hr = StartProcess(cmdline);

		if (FAILED(hr))
		{
			ReportError(hr);
		}

		return hr;
	}

	return S_FALSE;
}

void CDllContextMenu::CVerb::ReportError(HRESULT hr)
{
	TCHAR buffer[1024];

	DWORD cch = FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM,
               		          NULL, hr, 0, buffer, _countof(buffer), NULL);
	if (cch > 0)
	{
		USES_CONVERSION;
		MessageBox(NULL, buffer, W2CT(ProjectName), MB_OK|MB_ICONERROR);
	}
}

HRESULT CDllContextMenu::CVerb::StartProcess(const tstring& cmdline)
{
	vector<TCHAR> buff(cmdline.size());
	stdext::checked_copy(cmdline.begin(), cmdline.end(), buff.begin());

	STARTUPINFO si = { sizeof(STARTUPINFO) };
	PROCESS_INFORMATION pi = { 0, 0, 0, 0 };
	
	if (CreateProcess(NULL, &buff[0], 
		              NULL, NULL, FALSE, 0, NULL, NULL, &si, &pi))
	{
		CloseHandle(pi.hProcess);
		CloseHandle(pi.hThread);

		return S_OK;
	}

	return AtlHresultFromLastError();
}

tstring CDllContextMenu::CVerb::QueryStringValue(CRegKey& key, LPCTSTR name)
{
	TCHAR szValue[32768];
	ZeroMemory(szValue, sizeof(szValue));

	DWORD cchValue = _countof(szValue) - 1;

	if (key.QueryStringValue(name, szValue, &cchValue) == ERROR_SUCCESS)
	{
		return tstring(szValue, cchValue);
	}

	return tstring();
}


CDllContextMenu::CSelectedFile::CSelectedFile(const tstring &path)
	: m_Path(path)
	, m_Type(FileTypeInvalid)
{
	CPEImage image(path);

	if (image.IsManagedAssembly())
	{
		m_Type |= FileTypeManaged;
	}

	if (image.IsWin32PE())
	{
		m_Type |= FileTypeWin32;
	}

	if (image.IsComServer())
	{
		m_Type |= FileTypeCOM;
	}
}

